package com.hyzs.spark.sql

import org.apache.spark.ml.{Pipeline, PipelineModel, PipelineStage}
import org.apache.spark.ml.feature.{IndexToString, StringIndexer, StringIndexerModel, VectorAssembler}
import org.apache.spark.rdd.RDD
import org.apache.spark.mllib.linalg.{Vector, Vectors}
import org.apache.spark.mllib.regression.LabeledPoint
import org.apache.spark.mllib.util.MLUtils
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.functions._
import org.apache.spark.sql._
import org.apache.spark.sql.types._

import scala.collection.mutable.{ArrayBuffer, ListBuffer}
import com.hyzs.spark.utils.SparkUtils._
import com.hyzs.spark.utils.{BaseUtil, InferSchema, JsonUtil, Params, SparkUtils}
import java.math.BigDecimal

import com.fasterxml.jackson.databind.{DeserializationFeature, ObjectMapper}
import com.fasterxml.jackson.module.scala.DefaultScalaModule
import com.hyzs.spark.bean._

/**
  * Created by XIANGKUN on 2017/12/5.
  */

object ConvertLibsvm {


  val originalKey = "user_id"
  val key = "user_id_md5"
  val indexerArray = new ArrayBuffer[StringIndexerModel]
  val maxLabelLength = 1000

  def getIndexers(df: DataFrame, col: String): (String, StringIndexerModel) = {
    val indexer = new StringIndexer()
      .setInputCol(col)
      .setOutputCol(s"${col}_indexer")
      .setHandleInvalid("skip")
      .fit(df)
    (col,indexer)
  }

  def castStringType(df:DataFrame, col:String): (DataFrame, StringIndexerModel) = {
    val indexer = new StringIndexer()
      .setInputCol(col)
      .setOutputCol(s"${col}_indexer")
      .setHandleInvalid("skip")
      .fit(df)
    val transformed = indexer.transform(df)
    val res = transformed.withColumn(col, transformed(s"${col}_indexer")).drop(s"${col}_indexer")
    (res, indexer)
  }

  def replaceIndexedCols(df:DataFrame, cols:Seq[String]): DataFrame = {
    val remainCols = df.columns diff cols
    val replaceExprs = cols.map( col => s" ${col}_indexer as $col")
    df.selectExpr(remainCols ++: replaceExprs: _*)
  }

  def dropOldCols(df:DataFrame,
                  stringCols:Seq[String],
                  timeCols:Seq[String],
                  numberCols:Seq[String]): Option[DataFrame] = {
    // add string index start with 1
    val strExprs = stringCols.map(col => s" (${col}_indexer + 1) as $col")
    val timeExprs = timeCols.map(col => s" ${col}_stamp as $col")
    val numberExprs = numberCols.map(col =>  s" ${col}_number as $col")
    Some(df.selectExpr(strExprs ++: timeExprs ++: numberExprs :_*))
  }

  def castTimestampFuc(time:String): Long = {
    BaseUtil.getUnixStamp(time).getOrElse(0)
  }

  def castLibsvmString(label:String="0.0", row: Row): String = {
    val datum = row.toSeq
    val resString = new StringBuilder(label)
    datum.zipWithIndex.foreach{ case (field,i) =>
        if(field != 0.0){
          val digit = new BigDecimal(field.toString)
          resString += ' '
          resString ++= s"${i+1}:${digit.toPlainString}"
        }
    }
    resString.toString()
  }

  def saveRdd(rdd:RDD[String], savePath:String): Unit = {
    if(checkHDFileExist(savePath))
      dropHDFiles(savePath)
    rdd.saveAsTextFile(savePath)
  }

  def buildObjRdd(dataSchema:StructType,
                  indexerArray:Seq[(String,StringIndexerModel)]): RDD[String] = {
    val objList:ListBuffer[BaseObj] = new ListBuffer
    objList += Ob1(0, Params.NO_TYPE, key)
    val indexerMap: Map[String,Map[String,Int]] = indexerArray.map{ case (name,model) =>
      if(model.labels.length<maxLabelLength)
        (name, model.labels.zip(1 to model.labels.length).toMap)
      else{
        val labels = model.labels.slice(0, maxLabelLength)
        (name, labels.zip(1 to maxLabelLength).toMap)
      }
    }.toMap
    (1 to dataSchema.length).zip(dataSchema).foreach{ case (index, field) =>
      val obj = field.dataType match {
        case IntegerType => Ob1(index, Params.NUMERIC_TYPE, field.name)
        case DoubleType => Ob1(index, Params.NUMERIC_TYPE, field.name)
        case DateType => Ob1(index, Params.DATE_TYPE, field.name)
        case TimestampType => Ob1(index, Params.DATE_TYPE, field.name)
        case StringType => Ob2(index, Params.STRING_TYPE, field.name,
          indexerMap.getOrElse(field.name, Map("null"->0)))
        case _ => Ob1(index, Params.NUMERIC_TYPE, field.name)
      }
      objList += obj
    }
    val objRdd = sc.parallelize(objList).sortBy(obj => obj.key)
    val objStr = objRdd.mapPartitions( objs => {
      val mapper = new ObjectMapper
      mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)
      mapper.registerModule(DefaultScalaModule)
      objs.map(obj => mapper.writeValueAsString(obj))
    })
    objStr
  }

  def readObj(filePath:String): Array[Ob1] = {
    val objRdd = sc.textFile(filePath)
    val objList = objRdd.collect().map{ record =>
      broadMapper.value.registerModule(DefaultScalaModule)
      val obj = broadMapper.value.readValue(record, classOf[Ob1])
      obj
    }
    objList
  }

  def getStructToJson(dataSchema:StructType): RDD[String] = {
    val structArray = (1 to dataSchema.length).zip(dataSchema).map{ case (index,field) =>
      val info = StructInfo(index, field.name, field.dataType.typeName)
      broadMapper.value.registerModule(DefaultScalaModule)
      broadMapper.value.writeValueAsString(info)
    }
    sc.makeRDD(structArray)
  }

  def main(args: Array[String]): Unit = {

    val tableName = "hyzs.m1_train"
    val convertPath = "/hyzs/convert_data/"
    val taskPath = s"$convertPath$tableName/"
    val objPath = s"${taskPath}obj"
    val modelPath = s"${taskPath}model"
    val libsvmPath = s"${taskPath}libsvm"
    val indexPath = s"${taskPath}index"
    val namePath = s"${taskPath}name"
    val resultPath = s"${taskPath}result"

    val sourceData = sqlContext.table(tableName).drop(originalKey)
    val allLabel = sqlContext.table("hyzs.m1_label")
    // filter label based on source data
    val fullData = allLabel.join(sourceData, Seq(key), "right")

    val labelRdd:RDD[String] = fullData.select("label").rdd.map(row => row(0).toString)
    val indexRdd:RDD[String] = fullData.select(key).rdd.map(row => row(0).toString)
    val data = fullData.drop(key).drop("label").na.fill("0")
      .na.replace("*", Map("" -> "0", "null" -> "0"))
    val nameRdd = sc.makeRDD[String](sourceData.columns)

    val dataColsArray = data.columns
    var stringCols = Seq[String]()
    var timeCols = Seq[String]()
    var numberCols = Seq[String]()
    var result: Option[DataFrame] = None
    if(args.length >0 && args(0) == "test"){
      val pipeline = Pipeline.load(modelPath)
      result = Some(pipeline.fit(data).transform(data))

      val objList = readObj(s"$resultPath/$tableName.obj")
      stringCols = objList.filter(obj => obj.value == Params.STRING_TYPE).map(_.fieldName)
      timeCols = objList.filter(obj => obj.value == Params.DATE_TYPE).map(_.fieldName)
      numberCols = objList.filter(obj => obj.value == Params.NUMERIC_TYPE).map(_.fieldName)

    } else if(args.length >0 && args(0) == "train"){
      val dataSchema = InferSchema.inferSchema(data)
      val stringSchema = dataSchema.filter(field => field.dataType == StringType)
      val timeSchema = dataSchema.filter(field => field.dataType == TimestampType)

      val indexerArray = stringSchema.map(field => getIndexers(data, field.name))
      val objRdd = buildObjRdd(dataSchema, indexerArray)
      val pipeline = new Pipeline().setStages(Array(indexerArray.map(_._2): _*))
      stringCols = stringSchema.map(field => field.name)
      timeCols = timeSchema.map(field => field.name)
      numberCols = data.columns diff stringCols diff timeCols

      if(checkHDFileExist(modelPath))dropHDFiles(modelPath)
      pipeline.save(modelPath)

      result = Some(pipeline.fit(data).transform(data))

      saveRdd(objRdd, objPath)
      saveRdd(nameRdd, namePath)
      mkHDdir(resultPath)
      copyMergeHDFiles(s"$objPath/", s"$resultPath/$tableName.obj")
      copyMergeHDFiles(s"$namePath/", s"$resultPath/$tableName.name")
    }

    val stampUdf = udf(castTimestampFuc _)
    for(col <- timeCols){
      result = Some(result.get.withColumn(s"${col}_stamp", stampUdf(result.get(col))))
    }

    for(col <- numberCols){
      result = Some(result.get.withColumn(s"${col}_number", result.get(col).cast(DoubleType)))
    }
    result = dropOldCols(result.get, stringCols, timeCols, numberCols)
    result = Some(result.get.selectExpr(dataColsArray:_*))

    // zip rdd should have THE SAME partitions
    val libsvm = result.get.rdd.zip(labelRdd).map{
      case (row, i) => castLibsvmString(i, row)
    }

    saveRdd(libsvm, libsvmPath)
    saveRdd(indexRdd, indexPath)
    mkHDdir(resultPath)
    copyMergeHDFiles(s"$libsvmPath/", s"$resultPath/$tableName.libsvm")
    copyMergeHDFiles(s"$indexPath/", s"$resultPath/$tableName.index")
  }
}


